from docx import Document
from typing import List, Dict
import zipfile
from xml.etree import ElementTree
from .base_parser import BaseDocumentParser

class WordParser(BaseDocumentParser):
    """Word文档解析器"""
    
    def extract_text(self, file_path: str) -> str:
        """提取Word文档文本内容"""
        doc = Document(file_path)
        text_content = []
        image_count = {}  # 记录每个段落的图片数量
        current_page = 1  # Word没有页码概念，使用段落序号代替
        
        for para in doc.paragraphs:
            if para.text.strip():
                text_content.append(para.text.strip())
            
            # 检查段落中的图片
            for run in para.runs:
                if run._element.findall('.//w:drawing'):
                    if current_page not in image_count:
                        image_count[current_page] = 0
                    image_count[current_page] += 1
                    placeholder = f"[图片_{current_page}_{image_count[current_page]}]"
                    text_content.append(placeholder)
        
        return "\n".join(text_content)
    
    def extract_images(self, file_path: str) -> List[Dict]:
        """提取Word文档中的图片"""
        images = []
        
        try:
            with zipfile.ZipFile(file_path) as docx:
                # 解析word/document.xml获取图片引用
                xml_content = docx.read('word/document.xml')
                tree = ElementTree.fromstring(xml_content)
                
                # 命名空间
                ns = {'w': 'http://schemas.openxmlformats.org/wordprocessingml/2006/main',
                     'a': 'http://schemas.openxmlformats.org/drawingml/2006/main'}
                
                # 获取所有图片引用
                image_refs = []
                for img in tree.findall('.//a:blip', ns):
                    image_refs.append(img.get('{http://schemas.openxmlformats.org/officeDocument/2006/relationships}embed'))
                
                # 读取图片数据
                for idx, img_ref in enumerate(image_refs, 1):
                    try:
                        # 从relationships文件中获取图片路径
                        rels_xml = docx.read('word/_rels/document.xml.rels')
                        rels_tree = ElementTree.fromstring(rels_xml)
                        
                        for rel in rels_tree.findall('.//Relationship'):
                            if rel.get('Id') == img_ref:
                                image_path = rel.get('Target')
                                if image_path.startswith('/'):
                                    image_path = image_path[1:]
                                    
                                # 读取图片数据
                                image_data = docx.read(f'word/{image_path}')
                                image_format = image_path.split('.')[-1].lower()
                                
                                images.append({
                                    'image_data': image_data,
                                    'format': image_format,
                                    'page': 1,  # Word文档没有页码概念
                                    'index': idx
                                })
                    except Exception as e:
                        continue
                        
        except Exception as e:
            pass
            
        return images 